#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
彩票遗漏数计算脚本 V2 - 支持位置型彩种
从lottery_results表读取开奖数据，计算每期每个号码的遗漏值，存储到lottery_omission表
遗漏值从1开始（表示已经N期没开了）

支持的彩种：
- 双色球(ssq)、大乐透(dlt)、七乐彩(qlc)、快乐8(kl8): 组选遗漏
- 福彩3D(fc3d)、排列3(pl3): 百位、十位、个位、组选遗漏
- 排列5(pl5): 万位、千位、百位、十位、个位遗漏
- 七星彩(qxc): 第1-6位、第7位(蓝球)遗漏
"""

import sys
import os
import json
import pymysql
from datetime import datetime

# 添加项目根目录到路径
script_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, script_dir)

from crawler.config.settings import DATABASE_CONFIG

class OmissionCalculatorV2:
    def __init__(self):
        self.db_config = DATABASE_CONFIG
        self.connection = None

    def connect(self):
        """连接数据库"""
        self.connection = pymysql.connect(
            host=self.db_config['host'],
            port=self.db_config['port'],
            user=self.db_config['user'],
            password=self.db_config['password'],
            database=self.db_config['database'],
            charset='utf8mb4'
        )

    def get_lottery_config(self, lottery_type):
        """获取彩种配置"""
        configs = {
            # 非位置型彩种
            'ssq': {
                'type': 'non-position',
                'zones': {'red': (1, 33), 'blue': (1, 16)}
            },
            'dlt': {
                'type': 'non-position',
                'zones': {'red': (1, 35), 'blue': (1, 12)}
            },
            'qlc': {
                'type': 'non-position',
                'zones': {'red': (1, 30), 'blue': (1, 30)}
            },
            'kl8': {
                'type': 'non-position',
                'zones': {'red': (1, 80)}
            },
            # 位置型彩种
            'fc3d': {
                'type': 'position',
                'positions': {
                    'pos0': (0, 9),  # 百位
                    'pos1': (0, 9),  # 十位
                    'pos2': (0, 9)   # 个位
                },
                'group_selection': True  # 支持组选
            },
            'pl3': {
                'type': 'position',
                'positions': {
                    'pos0': (0, 9),  # 百位
                    'pos1': (0, 9),  # 十位
                    'pos2': (0, 9)   # 个位
                },
                'group_selection': True  # 支持组选
            },
            'pl5': {
                'type': 'position',
                'positions': {
                    'pos0': (0, 9),  # 万位
                    'pos1': (0, 9),  # 千位
                    'pos2': (0, 9),  # 百位
                    'pos3': (0, 9),  # 十位
                    'pos4': (0, 9)   # 个位
                }
            },
            'qxc': {
                'type': 'position',
                'positions': {
                    'pos0': (0, 9),  # 第1位
                    'pos1': (0, 9),  # 第2位
                    'pos2': (0, 9),  # 第3位
                    'pos3': (0, 9),  # 第4位
                    'pos4': (0, 9),  # 第5位
                    'pos5': (0, 9),  # 第6位
                    'pos6': (0, 9)   # 第7位（蓝球）
                }
            }
        }
        return configs.get(lottery_type)

    def parse_numbers(self, all_numbers):
        """解析开奖号码"""
        main_numbers = []
        special_numbers = []
        
        if all_numbers and isinstance(all_numbers, list):
            if len(all_numbers) > 0 and all_numbers[0]:
                main_numbers = [int(n) for n in all_numbers[0].split() if n.strip()]
            if len(all_numbers) > 1 and all_numbers[1] and all_numbers[1] != 'None':
                special_numbers = [int(n) for n in all_numbers[1].split() if n.strip()]
        
        return main_numbers, special_numbers

    def calculate_non_position_lottery(self, lottery_type, config):
        """计算非位置型彩种的遗漏（双色球、大乐透等）"""
        print(f"\n{'='*60}")
        print(f"开始计算 {lottery_type} 的遗漏数（非位置型）")
        print(f"{'='*60}")
        
        cursor = self.connection.cursor(pymysql.cursors.DictCursor)
        
        # 获取所有历史数据
        cursor.execute("""
            SELECT period, all_numbers
            FROM lottery_results
            WHERE lottery_type = %s
            ORDER BY period ASC
        """, (lottery_type,))
        
        results = cursor.fetchall()
        print(f"读取到 {len(results)} 条历史记录")
        
        if not results:
            print(f"警告: {lottery_type} 没有历史数据")
            return
        
        # 初始化遗漏计数器
        omission_counters = {}
        for zone, (min_num, max_num) in config['zones'].items():
            omission_counters[zone] = {num: 1 for num in range(min_num, max_num + 1)}
        
        # 准备批量插入的数据
        omission_records = []
        
        # 遍历所有历史数据计算遗漏
        for idx, result in enumerate(results):
            period = result['period']
            all_numbers = result['all_numbers']
            
            # 解析开奖号码
            if isinstance(all_numbers, str):
                all_numbers = json.loads(all_numbers)
            
            main_numbers, special_numbers = self.parse_numbers(all_numbers)
            
            # 处理红球/主号码
            if 'red' in config['zones']:
                min_num, max_num = config['zones']['red']
                for num in range(min_num, max_num + 1):
                    omission_records.append({
                        'lottery_type': lottery_type,
                        'period': period,
                        'number_zone': 'red',
                        'position': None,
                        'number': num,
                        'omission': omission_counters['red'][num]
                    })

                    if num in main_numbers:
                        omission_counters['red'][num] = 1
                    else:
                        omission_counters['red'][num] += 1

            # 处理蓝球/特别号码
            if 'blue' in config['zones']:
                min_num, max_num = config['zones']['blue']
                for num in range(min_num, max_num + 1):
                    omission_records.append({
                        'lottery_type': lottery_type,
                        'period': period,
                        'number_zone': 'blue',
                        'position': None,
                        'number': num,
                        'omission': omission_counters['blue'][num]
                    })

                    if num in special_numbers:
                        omission_counters['blue'][num] = 1
                    else:
                        omission_counters['blue'][num] += 1

            # 进度显示
            if (idx + 1) % 100 == 0 or idx == len(results) - 1:
                print(f"处理进度: {idx + 1}/{len(results)}")
        
        # 清空该彩种的旧数据
        cursor.execute("DELETE FROM lottery_omission WHERE lottery_type = %s", (lottery_type,))
        print(f"清空旧数据: {cursor.rowcount} 条")
        
        # 批量插入新数据
        if omission_records:
            print(f"准备插入 {len(omission_records)} 条遗漏记录...")
            
            batch_size = 1000
            for i in range(0, len(omission_records), batch_size):
                batch = omission_records[i:i + batch_size]
                
                cursor.executemany("""
                    INSERT INTO lottery_omission 
                    (lottery_type, period, number_zone, position, number, omission)
                    VALUES (%(lottery_type)s, %(period)s, %(number_zone)s, %(position)s, %(number)s, %(omission)s)
                """, batch)
                
                print(f"  插入批次 {i // batch_size + 1}: {len(batch)} 条")
            
            self.connection.commit()
            print(f"✅ {lottery_type} 遗漏数据计算完成！")
        
        cursor.close()

    def calculate_position_lottery(self, lottery_type, config):
        """计算位置型彩种的遗漏（福彩3D、排列3、排列5、七星彩）"""
        print(f"\n{'='*60}")
        print(f"开始计算 {lottery_type} 的遗漏数（位置型）")
        print(f"{'='*60}")
        
        cursor = self.connection.cursor(pymysql.cursors.DictCursor)
        
        # 获取所有历史数据
        cursor.execute("""
            SELECT period, all_numbers
            FROM lottery_results
            WHERE lottery_type = %s
            ORDER BY period ASC
        """, (lottery_type,))
        
        results = cursor.fetchall()
        print(f"读取到 {len(results)} 条历史记录")
        
        if not results:
            print(f"警告: {lottery_type} 没有历史数据")
            return
        
        # 初始化遗漏计数器（按位置）
        omission_counters = {}
        for pos_key, (min_num, max_num) in config['positions'].items():
            omission_counters[pos_key] = {num: 1 for num in range(min_num, max_num + 1)}
        
        # 如果支持组选，初始化组选遗漏计数器
        if config.get('group_selection'):
            omission_counters['group'] = {num: 1 for num in range(0, 10)}
        
        # 准备批量插入的数据
        omission_records = []
        
        # 遍历所有历史数据计算遗漏
        for idx, result in enumerate(results):
            period = result['period']
            all_numbers = result['all_numbers']
            
            # 解析开奖号码
            if isinstance(all_numbers, str):
                all_numbers = json.loads(all_numbers)
            
            main_numbers, special_numbers = self.parse_numbers(all_numbers)
            
            # 处理组选遗漏（如果支持）
            if config.get('group_selection'):
                unique_numbers = list(set(main_numbers[:3]))  # 前3位的唯一号码
                
                for num in range(0, 10):
                    omission_records.append({
                        'lottery_type': lottery_type,
                        'period': period,
                        'number_zone': 'red',
                        'position': None,  # 组选没有position
                        'number': num,
                        'omission': omission_counters['group'][num]
                    })
                    
                    if num in unique_numbers:
                        omission_counters['group'][num] = 1
                    else:
                        omission_counters['group'][num] += 1
            
            # 处理各个位置的遗漏
            for pos_key, (min_num, max_num) in config['positions'].items():
                pos_index = int(pos_key.replace('pos', ''))
                
                # 获取该位置的开奖号码
                drawn_number = None
                if lottery_type == 'qxc' and pos_key == 'pos6':
                    # 七星彩的pos6是蓝球，优先从special_numbers获取
                    if special_numbers:
                        drawn_number = special_numbers[0]
                    elif pos_index < len(main_numbers):
                        drawn_number = main_numbers[pos_index]
                elif pos_index < len(main_numbers):
                    drawn_number = main_numbers[pos_index]
                
                for num in range(min_num, max_num + 1):
                    # 确定number_zone
                    current_number_zone = 'red'
                    if lottery_type == 'qxc' and pos_key == 'pos6':
                        current_number_zone = 'blue'
                    
                    omission_records.append({
                        'lottery_type': lottery_type,
                        'period': period,
                        'number_zone': current_number_zone,
                        'position': pos_key,
                        'number': num,
                        'omission': omission_counters[pos_key][num]
                    })
                    
                    # 更新遗漏计数器
                    if drawn_number is not None and num == drawn_number:
                        omission_counters[pos_key][num] = 1
                    else:
                        omission_counters[pos_key][num] += 1
            
            # 进度显示
            if (idx + 1) % 100 == 0 or idx == len(results) - 1:
                print(f"处理进度: {idx + 1}/{len(results)}")
        
        # 清空该彩种的旧数据
        cursor.execute("DELETE FROM lottery_omission WHERE lottery_type = %s", (lottery_type,))
        print(f"清空旧数据: {cursor.rowcount} 条")
        
        # 批量插入新数据
        if omission_records:
            print(f"准备插入 {len(omission_records)} 条遗漏记录...")
            
            batch_size = 1000
            for i in range(0, len(omission_records), batch_size):
                batch = omission_records[i:i + batch_size]
                
                cursor.executemany("""
                    INSERT INTO lottery_omission 
                    (lottery_type, period, number_zone, position, number, omission)
                    VALUES (%(lottery_type)s, %(period)s, %(number_zone)s, %(position)s, %(number)s, %(omission)s)
                """, batch)
                
                print(f"  插入批次 {i // batch_size + 1}: {len(batch)} 条")
            
            self.connection.commit()
            print(f"✅ {lottery_type} 遗漏数据计算完成！")
        
        cursor.close()

    def calculate_lottery_omission(self, lottery_type):
        """计算某个彩种的遗漏数"""
        config = self.get_lottery_config(lottery_type)
        
        if not config:
            print(f"警告: 未知的彩种类型 {lottery_type}")
            return
        
        if config['type'] == 'non-position':
            self.calculate_non_position_lottery(lottery_type, config)
        else:
            self.calculate_position_lottery(lottery_type, config)

    def calculate_by_periods(self, lottery_type, periods=None):
        """
        增量计算指定期号的遗漏值

        参数:
            lottery_type: 彩种类型
            periods: 要计算的期号列表，如果为None则计算所有期号（全量计算）
        """
        config = self.get_lottery_config(lottery_type)

        if not config:
            print(f"警告: 未知的彩种类型 {lottery_type}")
            return

        # 如果没有指定期号，执行全量计算
        if periods is None:
            print(f"执行全量计算 {lottery_type}")
            if config['type'] == 'non-position':
                self.calculate_non_position_lottery(lottery_type, config)
            else:
                self.calculate_position_lottery(lottery_type, config)
            return

        # 增量计算指定期号
        print(f"\n{'='*60}")
        print(f"开始增量计算 {lottery_type} 的遗漏数（期号：{', '.join(periods)}）")
        print(f"{'='*60}")

        cursor = self.connection.cursor(pymysql.cursors.DictCursor)

        try:
            # 获取要计算的期号的历史数据（包含这些期号及之前所有期号）
            # 构建5位和7位格式的期号条件
            period_conditions = []
            period_params = [lottery_type]

            for period in periods:
                period_5 = str(period)[-5:]  # 确保是5位格式
                period_7 = '20' + period_5 if len(str(period)) == 5 else str(period)
                period_conditions.append("(period = %s OR period = %s)")
                period_params.extend([period_5, period_7])

            where_clause = " OR ".join(period_conditions)

            cursor.execute(f"""
                SELECT period, all_numbers
                FROM lottery_results
                WHERE lottery_type = %s
                AND ({where_clause})
                ORDER BY CAST(period AS UNSIGNED) ASC
            """, period_params)

            target_results = cursor.fetchall()
            if not target_results:
                print(f"警告: {lottery_type} 没有找到期号 {', '.join(periods)} 的数据")
                return

            # 获取这些期号之前的历史数据，用于计算遗漏基数
            earliest_period = min(result['period'] for result in target_results)
            cursor.execute("""
                SELECT period, all_numbers
                FROM lottery_results
                WHERE lottery_type = %s
                AND period < %s
                ORDER BY period ASC
            """, (lottery_type, earliest_period))

            history_results = cursor.fetchall()
            all_results = history_results + target_results

            print(f"读取到 {len(history_results)} 条历史记录 + {len(target_results)} 条目标记录")

            if not all_results:
                print(f"警告: {lottery_type} 没有历史数据")
                return

            # 计算遗漏基数（从历史记录开始）
            omission_counters = self._calculate_omission_base(lottery_type, config, history_results)

            # 准备批量插入的数据
            omission_records = []

            # 遍历目标期号计算遗漏
            for idx, result in enumerate(target_results):
                period = result['period']
                all_numbers = result['all_numbers']

                # 解析开奖号码
                if isinstance(all_numbers, str):
                    all_numbers = json.loads(all_numbers)

                main_numbers, special_numbers = self.parse_numbers(all_numbers)

                # 生成该期的遗漏记录
                period_records = self._generate_period_omission_records(
                    lottery_type, period, config, omission_counters, main_numbers, special_numbers
                )
                omission_records.extend(period_records)

                # 更新遗漏计数器
                self._update_omission_counters(lottery_type, config, omission_counters, main_numbers, special_numbers)

                print(f"处理期号 {period}: 生成 {len(period_records)} 条遗漏记录")

            # 删除这些期号的旧遗漏数据（如果存在）
            # 构建5位和7位格式的期号条件
            delete_conditions = []
            delete_params = [lottery_type]

            for period in periods:
                period_5 = str(period)[-5:]  # 确保是5位格式
                period_7 = '20' + period_5 if len(str(period)) == 5 else str(period)
                delete_conditions.append("(period = %s OR period = %s)")
                delete_params.extend([period_5, period_7])

            delete_where_clause = " OR ".join(delete_conditions)

            cursor.execute(f"""
                DELETE FROM lottery_omission
                WHERE lottery_type = %s AND ({delete_where_clause})
            """, delete_params)
            print(f"清空旧数据: {cursor.rowcount} 条")

            # 批量插入新数据
            if omission_records:
                print(f"准备插入 {len(omission_records)} 条遗漏记录...")

                batch_size = 1000
                for i in range(0, len(omission_records), batch_size):
                    batch = omission_records[i:i + batch_size]

                    cursor.executemany("""
                        INSERT INTO lottery_omission
                        (lottery_type, period, number_zone, position, number, omission)
                        VALUES (%(lottery_type)s, %(period)s, %(number_zone)s, %(position)s, %(number)s, %(omission)s)
                    """, batch)

                    print(f"  插入批次 {i // batch_size + 1}: {len(batch)} 条")

                self.connection.commit()
                print(f"✅ {lottery_type} 期号 {', '.join(periods)} 遗漏数据增量计算完成！")

        except Exception as e:
            print(f"❌ 增量计算 {lottery_type} 遗漏数失败: {e}")
            import traceback
            traceback.print_exc()
            self.connection.rollback()
        finally:
            cursor.close()

    def _calculate_omission_base(self, lottery_type, config, history_results):
        """根据历史记录计算遗漏基数"""
        omission_counters = {}

        # 初始化遗漏计数器
        if config['type'] == 'non-position':
            for zone, (min_num, max_num) in config['zones'].items():
                omission_counters[zone] = {num: 1 for num in range(min_num, max_num + 1)}
        else:
            for pos_key, (min_num, max_num) in config['positions'].items():
                omission_counters[pos_key] = {num: 1 for num in range(min_num, max_num + 1)}

            if config.get('group_selection'):
                omission_counters['group'] = {num: 1 for num in range(0, 10)}

        # 根据历史记录更新遗漏计数器
        for result in history_results:
            all_numbers = result['all_numbers']
            if isinstance(all_numbers, str):
                all_numbers = json.loads(all_numbers)

            main_numbers, special_numbers = self.parse_numbers(all_numbers)
            self._update_omission_counters(lottery_type, config, omission_counters, main_numbers, special_numbers)

        return omission_counters

    def _generate_period_omission_records(self, lottery_type, period, config, omission_counters, main_numbers, special_numbers):
        """为指定期号生成遗漏记录"""
        records = []

        if config['type'] == 'non-position':
            # 非位置型彩种
            if 'red' in config['zones']:
                min_num, max_num = config['zones']['red']
                for num in range(min_num, max_num + 1):
                    records.append({
                        'lottery_type': lottery_type,
                        'period': period,
                        'number_zone': 'red',
                        'position': None,
                        'number': num,
                        'omission': omission_counters['red'][num]
                    })

            if 'blue' in config['zones']:
                min_num, max_num = config['zones']['blue']
                for num in range(min_num, max_num + 1):
                    records.append({
                        'lottery_type': lottery_type,
                        'period': period,
                        'number_zone': 'blue',
                        'position': None,
                        'number': num,
                        'omission': omission_counters['blue'][num]
                    })
        else:
            # 位置型彩种
            if config.get('group_selection'):
                unique_numbers = list(set(main_numbers[:3]))
                for num in range(0, 10):
                    records.append({
                        'lottery_type': lottery_type,
                        'period': period,
                        'number_zone': 'red',
                        'position': None,
                        'number': num,
                        'omission': omission_counters['group'][num]
                    })

            for pos_key, (min_num, max_num) in config['positions'].items():
                pos_index = int(pos_key.replace('pos', ''))

                drawn_number = None
                if lottery_type == 'qxc' and pos_key == 'pos6':
                    if special_numbers:
                        drawn_number = special_numbers[0]
                    elif pos_index < len(main_numbers):
                        drawn_number = main_numbers[pos_index]
                elif pos_index < len(main_numbers):
                    drawn_number = main_numbers[pos_index]

                for num in range(min_num, max_num + 1):
                    current_number_zone = 'red'
                    if lottery_type == 'qxc' and pos_key == 'pos6':
                        current_number_zone = 'blue'

                    records.append({
                        'lottery_type': lottery_type,
                        'period': period,
                        'number_zone': current_number_zone,
                        'position': pos_key,
                        'number': num,
                        'omission': omission_counters[pos_key][num]
                    })

        return records

    def _update_omission_counters(self, lottery_type, config, omission_counters, main_numbers, special_numbers):
        """更新遗漏计数器"""
        if config['type'] == 'non-position':
            # 更新红球/主号码
            if 'red' in config['zones']:
                min_num, max_num = config['zones']['red']
                for num in range(min_num, max_num + 1):
                    if num in main_numbers:
                        omission_counters['red'][num] = 1
                    else:
                        omission_counters['red'][num] += 1

            # 更新蓝球/特别号码
            if 'blue' in config['zones']:
                min_num, max_num = config['zones']['blue']
                for num in range(min_num, max_num + 1):
                    if num in special_numbers:
                        omission_counters['blue'][num] = 1
                    else:
                        omission_counters['blue'][num] += 1
        else:
            # 更新组选遗漏（如果支持）
            if config.get('group_selection'):
                unique_numbers = list(set(main_numbers[:3]))
                for num in range(0, 10):
                    if num in unique_numbers:
                        omission_counters['group'][num] = 1
                    else:
                        omission_counters['group'][num] += 1

            # 更新各个位置的遗漏
            for pos_key, (min_num, max_num) in config['positions'].items():
                pos_index = int(pos_key.replace('pos', ''))

                drawn_number = None
                if lottery_type == 'qxc' and pos_key == 'pos6':
                    if special_numbers:
                        drawn_number = special_numbers[0]
                    elif pos_index < len(main_numbers):
                        drawn_number = main_numbers[pos_index]
                elif pos_index < len(main_numbers):
                    drawn_number = main_numbers[pos_index]

                for num in range(min_num, max_num + 1):
                    if drawn_number is not None and num == drawn_number:
                        omission_counters[pos_key][num] = 1
                    else:
                        omission_counters[pos_key][num] += 1

    def calculate_all(self, lottery_types=None):
        """计算所有彩种的遗漏数"""
        if lottery_types is None:
            lottery_types = ['ssq', 'dlt', 'fc3d', 'pl3', 'pl5', 'qlc', 'qxc', 'kl8']

        start_time = datetime.now()
        print(f"\n开始时间: {start_time.strftime('%Y-%m-%d %H:%M:%S')}")

        for lottery_type in lottery_types:
            try:
                self.calculate_lottery_omission(lottery_type)
            except Exception as e:
                print(f"❌ 计算 {lottery_type} 遗漏数失败: {e}")
                import traceback
                traceback.print_exc()

        end_time = datetime.now()
        duration = (end_time - start_time).total_seconds()

        print(f"\n{'='*60}")
        print(f"全部完成!")
        print(f"结束时间: {end_time.strftime('%Y-%m-%d %H:%M:%S')}")
        print(f"总耗时: {duration:.2f} 秒")
        print(f"{'='*60}\n")

    def close(self):
        """关闭数据库连接"""
        if self.connection:
            self.connection.close()

def main():
    import argparse
    
    parser = argparse.ArgumentParser(description='计算彩票遗漏数 V2')
    parser.add_argument('--lottery', type=str, help='指定彩种（ssq/dlt等），不指定则计算所有')
    parser.add_argument('--all', action='store_true', help='计算所有彩种')
    
    args = parser.parse_args()
    
    calculator = OmissionCalculatorV2()
    
    try:
        calculator.connect()
        
        if args.lottery:
            calculator.calculate_lottery_omission(args.lottery)
        else:
            calculator.calculate_all()
            
    except Exception as e:
        print(f"❌ 发生错误: {e}")
        import traceback
        traceback.print_exc()
    finally:
        calculator.close()

if __name__ == '__main__':
    main()

